from transformers import AutoProcessor, SeamlessM4Tv2Model
import torchaudio

model_name="facebook/seamless-m4t-v2-large"
processor = AutoProcessor.from_pretrained(model_name)
model = SeamlessM4Tv2Model.from_pretrained(model_name)

# # from text->array（不支持中文）
# #支持目标语言：arb（现代标准阿拉伯语）、ben（孟加拉语）、cat（加泰罗尼亚语）、ces（捷克语）、cmn（普通话）、cym（威尔士语）、dan（丹麦语）、
# # deu（德语）、eng（英语）、est（爱沙尼亚语）、fin（芬兰语）、fra（法语）、hin（印地语）、ind（印度尼西亚语）、ita（意大利语）、jpn（日语）、
# # kan（卡纳达语）、kor（韩语）、mlt（马耳他语）、nld（荷兰语）、pes（波斯语）、pol（波兰语）、por（葡萄牙语）、ron（罗马尼亚语）、rus（俄语）、
# # slk（斯洛伐克语）、spa（西班牙语）、swe（瑞典语）、swh（斯瓦希里语）、tam（泰米尔语）、tel（泰卢固语）、tgl（他加禄语）、tha（泰语）、
# # tur（土耳其语）、ukr（乌克兰语）、urd（乌尔都语）、uzn（乌兹别克语）、vie（越南语）
# input_text="Hello, my dog is cute"
# text_src_lang="eng"
# text_tgt_lang="jpn"
# text_inputs = processor(text =input_text , src_lang=text_src_lang, return_tensors="pt")
# audio_array_from_text = model.generate(**text_inputs, tgt_lang=text_tgt_lang)[0].cpu().numpy().squeeze()
# print(f"{input_text} -> {audio_array_from_text}")

# from audio->array（不支持中文）
audio_freq=16_000
audio_tgt_lang="jpn"
wav_file="input_audios/参会者，您好.wav"
audio, orig_freq =  torchaudio.load(wav_file)
audio =  torchaudio.functional.resample(audio, orig_freq=orig_freq, new_freq=audio_freq) # must be a 16 kHz waveform array
audio_inputs = processor(audios=audio, return_tensors="pt")
audio_array_from_audio = model.generate(**audio_inputs, tgt_lang=audio_tgt_lang)[0].cpu().numpy().squeeze()
print(f"{wav_file} -> {audio_array_from_audio}")